/**
 * 
 */
package us.codecraft.webmagic.scheduler;

import org.apache.commons.codec.Charsets;

import com.google.common.hash.Funnel;

import redis.clients.jedis.Jedis;
import redis.clients.jedis.JedisPool;
import redis.clients.jedis.JedisPoolConfig;
import us.codecraft.webmagic.Request;
import us.codecraft.webmagic.Task;
import us.codecraft.webmagic.scheduler.component.DuplicateRemover;
import us.codecraft.webmagic.utils.BloomFilterHelper;

/**
 * @author Administrator
 *
 */
public class RedisBloomFilterDuplicateRemover implements DuplicateRemover {

	protected JedisPool pool;

	private static final String COUNT_PREFIX = "count_";

	private static final String BLOOM_FILTER_DUPLICATE_KEY = "bloom_filter_duplicate";

	private BloomFilterHelper<String> bloomFilterHelper = new BloomFilterHelper<>(
			(Funnel<String>) (from, into) -> into.putString(from, Charsets.UTF_8).putString(from, Charsets.UTF_8), 
			200000000, 0.05);

	public RedisBloomFilterDuplicateRemover(String host) {
		this(new JedisPool(new JedisPoolConfig(), host));
	}

	public RedisBloomFilterDuplicateRemover(JedisPool pool) {
		this.pool = pool;
	}

	@Override
	public boolean isDuplicate(Request request, Task task) {
		Jedis jedis = pool.getResource();
		try {
			boolean duplicateRet = isExists(jedis, request.getUrl());
			if (!duplicateRet) {
				jedis.incr(getCountKey(task));
			}
			return duplicateRet;
		} finally {
			if (jedis != null)
				jedis.close();
		}
	}

	@Override
	public void resetDuplicateCheck(Task task) {
		Jedis jedis = pool.getResource();
		try {
			jedis.del(getCountKey(task), BLOOM_FILTER_DUPLICATE_KEY);
		} finally {
			if (jedis != null)
				jedis.close();
		}
	}

	@Override
	public int getTotalRequestsCount(Task task) {
		Jedis jedis = pool.getResource();
		try {
			String count = jedis.get(getCountKey(task));
			return count == null ? 0 : Integer.valueOf(count);
		} finally {
			if (jedis != null)
				jedis.close();
		}
	}
	
	protected String getCountKey(Task task) {
		return COUNT_PREFIX + task.getUUID();
	}

	public boolean isExists(String value) {
		try (Jedis jedis = pool.getResource()){
			return isExists(jedis, value);
		} catch(Exception ex) {
			throw ex;
		}
	}
	
	public void addBloomFilter(String value) {
		try (Jedis jedis = pool.getResource()){
			addBloomFilter(jedis, value);
		} catch(Exception ex) {
			throw ex;
		}
	}
	
	protected boolean includeByBloomFilter(Jedis jedis, String value) {
		boolean ret = isExists(jedis, value);
		// 添加值
		if (!ret) {
			addBloomFilter(jedis, value);
		}
		return ret;
	}
	
	private boolean isExists(Jedis jedis, String value) {
		int[] offset = bloomFilterHelper.murmurHashOffset(value);
		boolean ret = true;
		// 检查是否包含值
		for (int i : offset) {
			if (!jedis.getbit(BLOOM_FILTER_DUPLICATE_KEY, i)) {
				ret = false;
			}
		}
		return ret;
	}
	
	private void addBloomFilter(Jedis jedis, String value) {
		int[] offset = bloomFilterHelper.murmurHashOffset(value);
		for (int i : offset) {
			jedis.setbit(BLOOM_FILTER_DUPLICATE_KEY, Long.valueOf(i), true);
		}
	}
	
}
